from file_valid.generate import generate, compliant_judge
from file_valid.query_rewrite import query_rewrite
from file_valid.recall import recall
from file_valid.rerank import rerank


class FileValidPipeline:
    def __init__(self, trunks):
        self.trunks = trunks

    async def process(self):
        recall_ = []
        for trunk in self.trunks:
            text = trunk["text"]
            rewrite_result = await query_rewrite(text)
            recall_result = recall(rewrite_result)
            rerank_result = rerank(recall_result)
            if len(rerank_result) > 0:
                judge_result = await compliant_judge(text, [x.content for x in rerank_result])
                if judge_result:
                    recall_.append(await generate(rerank_result, text))
        if len(recall_) == 0:
            yield "您好！您上传的文件已通过系统基于知识库的扫描检测，暂未发现明确违规内容。"
        else:
            for x in recall_:
                async for chunk in x:
                    yield chunk
